import requests
import base64
import io
import json
import numpy as np
from PIL import Image
import time
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 生成测试图像
def generate_test_image(width=256, height=256, color='red'):
    """生成指定大小和颜色的测试图像"""
    image = Image.new('RGB', (width, height), color=color)
    
    # 转换为base64
    buffered = io.BytesIO()
    image.save(buffered, format="JPEG")
    image_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
    
    logger.info(f"生成测试图像: {width}x{height}, 颜色: {color}, base64长度: {len(image_base64)}")
    return image_base64

# 测试embedding生成和加载
def test_embedding_load(url="http://localhost:8077/embedding/base64", use_json=False):
    """测试embedding生成和加载
    Args:
        url: API接口地址
        use_json: 是否使用JSON格式提交请求
    """
    # 生成测试图像base64
    image_base64 = generate_test_image()
    
    # 准备请求
    if use_json:
        headers = {'Content-Type': 'application/json'}
        data = json.dumps({'image_base64': image_base64})
        logger.info("使用JSON格式提交请求")
    else:
        headers = {'Content-Type': 'application/x-www-form-urlencoded'}
        data = {'image_base64': image_base64}
        logger.info("使用表单格式提交请求")
    
    try:
        logger.info(f"发送请求到 {url}...")
        start_time = time.time()
        response = requests.post(url, headers=headers, data=data)
        end_time = time.time()
        
        logger.info(f"请求完成，耗时: {end_time - start_time:.2f}秒，状态码: {response.status_code}")
        
        if response.status_code == 200:
            try:
                result = response.json()
                logger.info(f"返回结果包含字段: {result.keys()}")
                logger.info(f"embedding形状: {result['shape']}")
                logger.info(f"image_id: {result['image_id']}")
                logger.info(f"model_type: {result['model_type']}")
                
                # 检查embedding数据
                if 'embedding' not in result:
                    logger.error("返回结果中缺少'embedding'字段")
                    return False
                
                # 尝试加载embedding数据
                try:
                    # 检查数据类型
                    logger.info(f"embedding数据类型: {type(result['embedding'])}")
                    
                    # 检查第一级元素数量
                    logger.info(f"embedding第一级元素数量: {len(result['embedding'])}")
                    
                    # 转换为numpy数组
                    embedding = np.array(result['embedding'])
                    logger.info(f"成功加载embedding为numpy数组，形状: {embedding.shape}")
                    
                    # 检查数组值范围
                    logger.info(f"embedding值范围: [{np.min(embedding):.4f}, {np.max(embedding):.4f}]")
                    
                    # 检查是否包含NaN或无穷大值
                    if np.isnan(embedding).any():
                        logger.warning("embedding中包含NaN值")
                    if np.isinf(embedding).any():
                        logger.warning("embedding中包含无穷大值")
                    
                    return True
                except Exception as e:
                    logger.error(f"加载embedding失败: {e}")
                    # 打印原始数据的前100个字符
                    raw_data = str(result['embedding'])[:100] + '...'
                    logger.error(f"原始embedding数据示例: {raw_data}")
                    return False
            except json.JSONDecodeError as e:
                logger.error(f"解析JSON响应失败: {e}")
                logger.error(f"响应内容: {response.text[:100]}...")
                return False
        else:
            logger.error(f"请求失败，状态码: {response.status_code}")
            logger.error(f"错误信息: {response.text[:200]}...")
            return False
    except Exception as e:
        logger.error(f"请求发生异常: {e}")
        return False

# 测试不同图像大小
def test_different_image_sizes():
    """测试不同大小的图像对embedding生成的影响"""
    sizes = [(128, 128), (256, 256), (512, 512), (1024, 1024)]
    results = {}
    
    for width, height in sizes:
        logger.info(f"测试图像大小: {width}x{height}")
        image_base64 = generate_test_image(width=width, height=height)
        
        url = "http://localhost:8077/embedding/base64"
        headers = {'Content-Type': 'application/x-www-form-urlencoded'}
        data = {'image_base64': image_base64}
        
        try:
            response = requests.post(url, headers=headers, data=data)
            if response.status_code == 200:
                result = response.json()
                embedding_shape = result['shape']
                logger.info(f"图像大小 {width}x{height}，生成的embedding形状: {embedding_shape}")
                results[(width, height)] = embedding_shape
            else:
                logger.error(f"测试图像大小 {width}x{height} 失败，状态码: {response.status_code}")
                results[(width, height)] = None
        except Exception as e:
            logger.error(f"测试图像大小 {width}x{height} 发生异常: {e}")
            results[(width, height)] = None
    
    return results

# 主函数
def main():
    logger.info("开始embedding详细测试...")
    
    # 测试表单格式提交
    logger.info("=== 测试1: 使用表单格式提交请求 ===")
    form_success = test_embedding_load(use_json=False)
    
    # 测试JSON格式提交
    logger.info("=== 测试2: 使用JSON格式提交请求 ===")
    json_success = test_embedding_load(use_json=True)
    
    # 测试不同图像大小
    logger.info("=== 测试3: 测试不同图像大小 ===")
    size_results = test_different_image_sizes()
    
    # 输出总结
    logger.info("=== 测试总结 ===")
    logger.info(f"表单格式提交测试: {'成功' if form_success else '失败'}")
    logger.info(f"JSON格式提交测试: {'成功' if json_success else '失败'}")
    
    logger.info("不同图像大小测试结果:")
    for size, shape in size_results.items():
        logger.info(f"  图像大小 {size[0]}x{size[1]}: {shape if shape else '失败'}")
    
    if form_success:
        logger.info("测试完成，服务端能够正确生成和返回embedding数据。")
        logger.info("如果您在自己的应用中遇到加载问题，可能是客户端代码或数据处理问题。")
    else:
        logger.info("测试失败，请检查服务端配置和运行状态。")

if __name__ == "__main__":
    main()

# 常见问题排查建议:
# 1. 确保客户端请求格式正确，使用表单格式(application/x-www-form-urlencoded)提交image_base64字段
# 2. 检查网络连接和服务可用性
# 3. 确认返回的JSON数据结构正确，包含embedding字段
# 4. 在客户端使用try-except捕获可能的异常，并打印详细错误信息
# 5. 如果使用前端框架，确保正确处理大数和嵌套数组